import matplotlib.pyplot as plt
import numpy as np
import rasterio
import os
from matplotlib import rcParams
from rasterio.mask import mask
import matplotlib.font_manager as fm
import geopandas as gpd
import os
pwd = os.getcwd()
"""
nc转换为tiff时候，查看scale_factor和offset不存在，不对此做处理
"""
# 设置中文字体
rcParams['font.sans-serif'] = ['Noto Sans CJK JP']  # 选择简体中文
rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 定义文件夹路径
rainfall_path = './mean_rain_tif/'
temperature_path = './mean_tem_tif/'

# 初始化存储结果的列表
rainfall_mean = []
temperature_mean = []

mask_name = 'shanxi'
mask_chinese = '陕西'
shapefile_path = '{}.shp'.format(mask_name)
mask_gdf = gpd.read_file(os.path.join(pwd,'shanxi_mask',shapefile_path))


# 读取降雨数据并计算平均值
for file in sorted(os.listdir(rainfall_path)):
    with rasterio.open(os.path.join(rainfall_path, file)) as src:
        nodata = 0
        out_image, out_transform = mask(src, mask_gdf.geometry, crop=True)
        image = out_image[0]  # 读取第一个波段
        image_except_nodata = image[image != nodata]
        rainfall_mean.append(image_except_nodata.mean())

# 读取气温数据并计算平均值
for file in sorted(os.listdir(temperature_path)):
    with rasterio.open(os.path.join(temperature_path, file)) as src:
        nodata = src.nodata
        out_image, out_transform = mask(src, mask_gdf.geometry, crop=True)
        image = out_image[0]  #
        # image = src.read(1)  # 读取第一个波段
        image_except_nodata = image[image != nodata]
        temperature_mean.append(image_except_nodata.mean())

# 转换为 numpy 数组
rainfall_mean = np.array(rainfall_mean)
temperature_mean = np.array(temperature_mean)

# 定义年份
years = np.arange(2002, 2002+ len(rainfall_mean))

# 创建图表
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(10, 8))

# 上图：年降雨量和趋势线
ax1.plot(years, rainfall_mean, marker='s', label='降雨量', color='black')
z_rainfall = np.polyfit(years, rainfall_mean, 1)
p_rainfall = np.poly1d(z_rainfall)
ax1.plot(years, p_rainfall(years), "r-", label='趋势线')
ax1.axhline(y=np.mean(rainfall_mean), color='gray', linestyle='--', label='平均值')  # 添加灰色水平线
ax1.set_ylabel('陕西年平均降雨量 (mm)')
ax1.legend()
ax1.text(0.25, 0.95, f'y={z_rainfall[0]:.4f}x+{z_rainfall[1]:.2f}\n$R^2$={np.corrcoef(years, rainfall_mean)[0, 1]**2:.8f}', 
         transform=ax1.transAxes, fontsize=12, verticalalignment='top')
ax1.set_xticks(years)  # 设置年份刻度
ax1.set_xticklabels(years, rotation=45)
# 下图：年气温和趋势线
ax2.plot(years, temperature_mean, marker='s', label='气温', color='black')
z_temperature = np.polyfit(years, temperature_mean, 1)
p_temperature = np.poly1d(z_temperature)
ax2.plot(years, p_temperature(years), "r-", label='趋势线')
ax2.axhline(y=np.mean(temperature_mean), color='gray', linestyle='--', label='平均值')  # 添加灰色水平线
ax2.set_ylabel('陕西年平均气温 (°C)')
ax2.set_xlabel('年份')
ax2.legend()
ax2.text(0.05, 0.95, f'y={z_temperature[0]:.4f}x+{z_temperature[1]:.2f}\n$R^2$={np.corrcoef(years, temperature_mean)[0, 1]**2:.8f}', 
         transform=ax2.transAxes, fontsize=12, verticalalignment='top')
ax2.set_xticks(years)  # 设置年份刻度
ax2.set_xticklabels(years, rotation=45)
plt.tight_layout()
plt.savefig('rainfall_temperature_trend.png', dpi=300)
plt.savefig('mean_rain_tem.png', dpi=300)
plt.show()